#!/usr/bin/perl -w

use strict;
use Getopt::Std;
use Net::SNMP;

#################
# Sub prototypes
################
sub get_iface_traffic;
sub HELP_MESSAGE;
sub set_iface_down;
sub setup;
sub sigint_handler;
sub snmp_connect;
sub snmp_get;
sub snmp_set;
sub start;
sub VERSION_MESSAGE;

############
# Constants
###########
$Getopt::Std::STANDARD_HELP_VERSION = 1;

############
# Variables
###########
my $version = 0.1;
my $configFile = 'conf.nq';
my $options = {};
my $quit;

my $session;
my @ifaces = ();
my $oid = {};
$oid->{'name'} = '1.3.6.1.2.1.2.2.1.2.';
$oid->{'speed'} = '1.3.6.1.2.1.2.2.1.5.';
$oid->{'status'} = '1.3.6.1.2.1.2.2.1.7.';
$oid->{'in'} = '1.3.6.1.2.1.2.2.1.10.';
$oid->{'out'} = '1.3.6.1.2.1.2.2.1.16.';

#############
# Main entry
############
getopts 'f:hsv', $options;

#############
# Read flags
############
$configFile = $options->{'f'} if (defined $options->{'f'});
if (defined $options->{'h'}) {
    &VERSION_MESSAGE();
    &HELP_MESSAGE();
    exit 0;
}
if (defined $options->{'v'}) {
    &VERSION_MESSAGE();
    exit 0;
}

##################
# Connect to SNMP
#################
print "> Connecting to SNMP agent...\n";
$session = &snmp_connect('172.19.13.12', 161, 'private');
if (!defined $session) {
    print "! Error during SNMP connection !\n";
    exit 1;
}

#######################
# Shall it run setup ?
######################
&setup() if (defined $options->{'s'});

########################
# Start the application
#######################
&start();

exit 0;

#######
# Subs
######

##
# get_iface_trafic
# get in/out/total octets
# param $_[0] interface
#
sub get_iface_traffic {
    my $session = shift;
    my $iface = shift;
    
    $iface->{'in'} = &snmp_get($session, $oid->{'in'} . $iface->{'id'}) / 1048576;
    if (!defined $iface->{'in'}) {
        print "! Error during SNMP get !\n";
        return;
    }
    $iface->{'out'} = &snmp_get($session, $oid->{'out'} . $iface->{'id'}) / 1048576;
    if (!defined $iface->{'out'}) {
        print "! Error during SNMP get !\n";
        return;
    }
    $iface->{'total'} = $iface->{'in'} + $iface->{'out'};
}

##
# HELP_MESSAGE
# Display help
#
sub HELP_MESSAGE {
    print "Usage: $0 [OPTIONS]...\n\n";
    print "Valid options:\n";
    print "  -f              configuration file to use (default: $configFile)\n";
    print "  -h, --help      this help screen\n";
    print "  -s              setup mode\n";
    print "  -v, --version   display version and exit\n";
    print "\nExit status is 0 if OK, 1 if error.\n";
}

##
# set_iface_down
# param $_[0] interface
#
sub set_iface_down {
    my $session = shift;
    my $iface = shift;
    
    $iface->{'status'} = 2;
    &snmp_set($session, $oid->{'status'} . $iface->{'id'}, INTEGER, 2);
}

sub setup {
}

##
# sigint_handler
# interrupt signal handler
#
sub sigint_handler {
    print "> Received SIGINT! Terminating...\n";
    $quit = 1;
}

##
# snmp_connect
#
sub snmp_connect {
    my $host = shift;
    my $port = shift;
    my $community = shift;
    
    my ($session, $error) = Net::SNMP->session(-hostname => $host, -port => $port, -community => $community, -version => 'snmpv2c');
    
    if (!defined $session) {
        return;
    }
    else {
        return $session;
    }
}

##
# snmp_get
#
sub snmp_get {
    my $session = shift;
    my $oid = shift;
    
    my $result = $session->get_request(-varbindlist => [$oid]);
    
    if (!defined $result) {
        return;
    }
    else {
        return $result->{$oid};
    }
}

##
# snmp_set
#
sub snmp_set {
    my $session = shift;
    my $oid = shift;
    my $type = shift;
    my $value = shift;
    
    my $result = $session->set_request(-varbindlist => [$oid, $type, $value]);
    
    if (!defined $result) {
        return;
    }
    else {
        return $result->{$oid};
    }
}

##
# start
# Main process
#
sub start {
    my $fd;
    my $i = 0;
    my $pattern = '^([0-9]+):(.+):(.+):([0-9\.]+)$';
    $SIG{INT} = \&sigint_handler;
    
    ##########################
    # Does conf file exists ?
    #########################
    if (!-e $configFile) {
        print "! Could not find configuration file ! Please check that $configFile exists !\n";
        exit 1;
    }
    ###################
    # Is it readable ?
    ##################
    if (!-r $configFile) {
        print "! Found configuration file, but you don't have rights to read it !\n";
        exit 1;
    }
    #########################
    # Does it contain data ?
    ########################
    if (-z $configFile) {
        print "! Configuration file seems to be empty ! Please run NetQuota with the -s flag !\n";
        exit 1;
    }
    ################################
    # File is OK, let's read it ...
    ###############################
    open $fd, $configFile or die "!! Could not open $configFile !\n";
    print "> Reading configuration file ...\n";
    while (<$fd>) {
        if (/$pattern/) {
            my $iface = {};
            $iface->{'id'}    = $1;
            $iface->{'name'}  = $2;
            $iface->{'descr'} = $3;
            $iface->{'quota'} = $4;
            $ifaces[$i] = $iface;
            $i++;
        }
    }
    close $fd;
    
    ###########################
    # Display found interfaces
    ##########################
    foreach my $iface (@ifaces) {
        print "#$iface->{'id'} $iface->{'name'} $iface->{'descr'} $iface->{'quota'}\n";
    }
    
    ############
    # Main loop
    ###########
    while (!$quit) {
        foreach my $iface (@ifaces) {
            # If the iface is up, we check its traffic
            #if ($iface->{'status'} == 1) {
                get_iface_traffic($session, $iface);
                # If the quota has been exceeded, we shut the iface down
                if ($iface->{'total'} > $iface->{'quota'})
                {
                    print "> Setting interface $iface->{'name'} down... ($iface->{'quota'} > $iface->{'total'})\n";
                    set_iface_down($session, $iface);
                }
            #}
        }
        sleep 3;
    }
}

##
# VERSION_MESSAGE
# Display current version
#
sub VERSION_MESSAGE {
    print "NetQuota version $version\n";
    print "Copyright (C) 2008 Frederic Mereu\n";
    print "Licence GPLv2 GNU GPL version 2\n\n";
}

